import sys,os,pickle,numpy as np
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
import requests,time,json,glob,re
def open_pickle_file(file_path: str) -> dict:
    with open(file_path, 'rb') as f:
        return pickle.load(f)
def keypoint_selector_2(inputs):
        points = inputs[0]
        confidences = inputs[1]
        if len(points) == 2:
            head, tail = points
            head = head[0]
            tail = tail[0]
            center = (np.array(head) + np.array(tail)) / 2
            dy = head[1] - tail[1]
            dx = head[0] - tail[0]
            theta = np.degrees(np.arctan2(dy, dx))
            return [[center[0].item(), center[1].item(), theta.item()], np.mean(confidences).item()]
        else:
            raise NotImplementedError
def sync_req(base, task, inputs):
    js = requests.post(base + '/task/' + task, json=inputs).json()
    tid = js["task_id"]
    prog = base + js["progress_pipe_addr"]
    res = base + js["result_addr"]
    print(f'start task {task} id {tid}', end='')
    while True:
        if requests.get(prog).json()["status"] == "Completed":
            break
        time.sleep(1)
        print('.', end='')
    res = requests.get(res).json()
    print()
    return res
def instant_req(url, inputs):
    resp = requests.post(url, json=inputs).json()
    return resp["instant_result"]
from protocol.decorators import edp_service, instant
@edp_service
class PipeLine:
    name = "py-pipeline"
    desc = "PIPELINE"
    tasks = [
        'get_videos: (Str, Str) -> List<(Str, Str)>',
        'get_dlc_data: (Str, Str)->List<(PoseFlat, Confidence)>',
        'get_stim: (Str, Bool, ?Bool)->(List<Int>, List<Int>)',
        'deal_stim: (List<Int>, List<Int>)->(List<(Int, Int)>, List<(Int, Int)>)',
        'deal_pose: List<(PoseFlat, Confidence)>->(List<(PoseFlat, Confidence)>, List<Float>)',
        'gen_data: (Str, Str, Bool, ?Bool, Int)->?List<List<(PoseFlat, DistributeStimulus)>>',
        'configure_init: (Str,Str)->Void',
        "configure_run: Str->CockroachDataset",
        "config_get_freqs: Str->Map<Str, Int>",
        "data_start: (Str,Str,Bool,Bool,Int)->((Str,Str),(Str,Bool,Bool), Int, Str)",
        "configure_xlsx: Str -> List<(Str,Str,Bool,Bool,Int)>",
        "merge: (([(PoseFlat, Confidence)], [Float]), ([(Int, Int)], [(Int, Int)]), Int, Str)->([(PoseFlat, Confidence)], [(Int, Int)], [(Int, Int)], Frequency, Str)"
    ]
    @instant
    def configure_xlsx(self, inputs):
        prj = inputs
        import pandas as pd
        df = pd.read_excel(prj + '/videos.xlsx')
        res = []
        for i in range(len(df)):
            each = df.iloc[i]
            _, video, fn, lr, freq = each.fillna('?')
            if lr == '?' or freq == '?':
                print(f'incomplete info: {video}, pass')
            elif os.path.exists(r"C:\Users\songy\Desktop\data-0530" + f"/{video}/ds.npz"):
                print(f"video already done: {video}")
            else:
                if lr == 2:
                    duel = True
                    lr = False
                elif lr == 1:
                    duel = False
                    lr = True
                elif lr == 0:
                    duel = False
                    lr = False
                else:
                    print(f'invalid lr: {video}')
                    continue
                freq = int(freq)
                res.append([prj, fn, duel, lr, freq])
        return res
    @instant
    def data_start(self, inputs):
        prj = inputs[0]
        fullname = inputs[1]
        video = fullname.rsplit('/', 1)[0].rsplit('\\', 1)[1].split('.', 1)[0]
        duel = inputs[2]
        lr = inputs[3]
        freq = inputs[4]
        saved = r"C:\Users\songy\Desktop\data-0530" + f"\{video}"
        return [[prj, fullname], [fullname,duel, lr], freq, saved]
    @instant
    def get_videos(self, inputs):
        vd = inputs[0]
        fregex = inputs[1]
        res = []
        for each in os.listdir(vd):
            if re.match(fregex, each):
                res.append([each.rsplit('.', 1)[0], vd + '/' + each])
        return res
    @instant
    def get_dlc_data(self, inputs):
        prj = inputs[0]
        video = inputs[1]
        import glob
        pkl = glob.glob(f'{prj}/videos/{video}*_full.pickle')[0]
        content = open_pickle_file(pkl)
        for each in content:
            if each == "metadata":
                continue
            coords = content[each]['coordinates'][0]
            confs = content[each]['confidence']
            content[each] = keypoint_selector_2((coords, confs))
        del content["metadata"]
        return [video, content]
    @instant
    def get_stim(self, inputs):
        video = inputs[0]
        duel = inputs[1]
        lr = inputs[2]
        thresh = 180
        lstim = sync_req("http://127.0.0.1:5020", "extract_stimulus", [video, thresh])
        if duel:
            rstim = sync_req("http://127.0.0.1:5020", "extract_stimulus", [video, thresh])
            return [lstim, rstim]
        if lr:
            return [[], lstim]
        else:
            return [lstim, []]
    @instant
    def deal_stim(self, inputs):
        lstim, rstim = inputs
        url = "http://127.0.0.1:5030" + "/task/stimulus_segment"
        sl = instant_req(url, [[x, False] for x in lstim])
        sr = instant_req(url, [[x, True] for x in rstim])
        return [sl, sr]
    @instant
    def deal_pose(self, inputs):
        thetas = [x[0][2] for x in inputs.values()]
        url = "http://127.0.0.1:5030" + "/task/normalize_thetas"
        smooth = instant_req(url, thetas)
        res = [[[x[0][0], x[0][1], y], x[1]] for (x, y) in zip(inputs.values(), thetas)]
        return res, smooth
    def regulate_pose0(self, inputs):
        return list(inputs.values())
    @instant
    def gen_data(self, inputs):
        video = inputs[0]
        fullpath = inputs[1]
        duel = inputs[2]
        lr = inputs[3]
        freq = inputs[4]
        prj = inputs[5]
        saved = r"C:\Users\songy\Desktop\data-0530" + f"\{video}"
        os.makedirs(saved, exist_ok=True)
        if not os.path.exists(saved+"/pose0.json"):
            _, content = self.get_dlc_data((prj, video))
            json.dump(content,open(saved+"/pose0.json","w"))
        else:
            content = json.load(open(saved+"/pose0.json","r"))
        if not os.path.exists(saved+"/lstim.json"):
            lstim, rstim = self.get_stim((fullpath, duel, lr))
            json.dump(lstim,open(saved+"/lstim.json","w"))
            json.dump(rstim,open(saved+"/rstim.json","w"))
        else:
            lstim = json.load(open(saved+"/lstim.json","r"))
            rstim = json.load(open(saved+"/rstim.json","r"))
        if not os.path.exists(saved+"/lstim_seg.json"):
            sl, sr = self.deal_stim((lstim, rstim))
            json.dump(sl,open(saved+"/lstim_seg.json","w"))
            json.dump(sr,open(saved+"/rstim_seg.json","w"))
        else:
            sl = json.load(open(saved+"/lstim_seg.json","r"))
            sr = json.load(open(saved+"/rstim_seg.json","r"))
        # pose0_regulate = self.regulate_pose0(content)
        # instant_req("http://127.0.0.1:3010" + "/task/pose_vis", pose0_regulate)
        if not os.path.exists(saved+"/pose.json"):
            res, smooth = self.deal_pose(content)
            json.dump(smooth,open(saved+"/smooth.json","w"))
            json.dump(res,open(saved+"/pose.json","w"))
        else:
            smooth = json.load(open(saved+"/smooth.json","r"))
            res = json.load(open(saved+"/pose.json","r"))
        if lstim or rstim:
            stim = sl[0] if lstim else sr[0]
            stim = [stim[0], stim[1], freq]
            instant_req("http://127.0.0.1:3010" + "/task/pose_single_stim", [res, stim])
        npz = saved + "/ds.npz"
        if not os.path.exists(npz):
            url = "http://127.0.0.1:5050" + "/task/generate_dataset_lrstim"
            resp = instant_req(url, (res, sl, sr, freq, npz))
            if not resp:
                print(f'error generating dataset: {video}')
                return None
        print(f'done: {video}')
        return None
    @instant
    def configure_init(self, inputs):
        vd = inputs[0]
        regex = inputs[1]
        import pandas as pd
        df = pd.DataFrame(columns=['video', 'fullname', 'lr', 'freq'])
        for video, fullname in self.get_videos((vd, regex)):
            df.loc[len(df)] = [video, fullname, None, None]
        df.to_excel(prj + '/videos.xlsx')
    @instant
    def configure_run(self, inputs):
        prj = inputs
        import pandas as pd
        df = pd.read_excel(prj + '/videos.xlsx')
        for i in range(len(df)):
            each = df.iloc[i]
            _, video, fn, lr, freq = each.fillna('?')
            fn = os.path.join(prj, 'videos', video + '.mp4')
            if lr == '?' or freq == '?':
                print(f'incomplete info: {video}, pass')
            # elif os.path.exists(r"C:\Users\songy\Desktop\data-0530" + f"/{video}/ds.npz"):
                # print(f"video already done: {video}")
            else:
                if lr == 2:
                    duel = True
                    lr = False
                elif lr == 1:
                    duel = False
                    lr = True
                elif lr == 0:
                    duel = False
                    lr = False
                else:
                    print(f'invalid lr: {video}')
                    continue
                freq = int(freq)
                self.gen_data((video, fn, duel, lr, freq, prj))
        return None
    @instant
    def config_get_freqs(self, inputs):
        prj = inputs
        import pandas as pd
        df = pd.read_excel(prj + '/videos.xlsx')
        res = {}
        for i in range(len(df)):
            each = df.iloc[i]
            _, video, fn, lr, freq = each.fillna('?')
            if freq == '?':
                print(f'incomplete info: {video}, pass')
            else:
                res[str(video)] = int(freq)
        return res

if __name__ == '__main__':
    prj = r"M:\fyp-songy-2025-05-17"
    vd = prj + '/videos'    
    fregex = r"fyp_\d+\.mp4"
    PipeLine.app.run(port=5060, debug=True)
    # PipeLine().configure_run(prj)
    # PipeLine().configure_init((vd, fregex))
    
    # video, fullname = PipeLine().get_videos(vd)[0]
    # duel = True
    # lr = None
    # freq = 50
    # PipeLine().gen_data((video, fullname, duel, lr, freq))